


import os 


import h5py    
import numpy as np    
import pdb

import debugpy

import matplotlib.pyplot as plt

import argparse
from os.path import join, split, exists,isdir

argsparser = argparse.ArgumentParser()

argsparser.add_argument('--src-path', required=True, default='',type=str)
argsparser.add_argument('--out-path', default='logs/',type=str)


args = argsparser.parse_args()



def load_hdf5(file_name):
    data = {}
    with h5py.File(file_name, "r") as hdf:
        # List the datasets inside the file
        # print("Datasets:", list(hdf.keys()))

        for key,value in hdf.items():
            if sum(ele for ele in value.shape) < 5 or "version" in key:
                continue
                # if value.dtype == "|S5":
                #     res.append((key, str(value).replace("[", "").replace("]", "").replace("b'", "").replace("'", "")))
                # else:
                #     res.append((key, value))
            else:
                # res.append((key, value.shape))
                data[key] = np.array(value)

    return data 
        
        


if __name__ == "__main__":

    # debugpy.listen(5678)
    # debugpy.wait_for_client()

    
    data = load_hdf5(args.src_path)

    if not exists(args.out_path):
        os.makedirs(args.out_path)

    
    for k,v in data.items():
        # plt.imshow(v, cmap="gray")  # Use 'gray' for grayscale images
        
        plt.imshow(v)  # Use 'gray' for grayscale images
        plt.axis("off")

        # pdb.set_trace()
        tmp_name = args.src_path.split('/')[-1].split('.')[-2]
        

        plt.savefig(join(args.out_path,'%s_%s.jpg'%(k, tmp_name)))

